import json
import sys
import os
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import networkx as nx
import yaml
from typing import Dict, Any

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.model import OpinionModel

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

def load_config(config_path: str) -> dict:
    """加载配置文件"""
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

def run_simulation(config: dict) -> dict:
    """运行模拟"""
    model = OpinionModel(config)
    
    # 记录模拟结果
    results = {
        "opinions": [],
        "group_opinions": [[] for _ in range(config["simulation"]["num_groups"])]
    }
    
    for step in range(config["simulation"]["num_iterations"]):
        state = model.step()
        
        # 记录所有智能体的观点
        opinions = [agent["opinion"] for agent in state["agents"]]
        results["opinions"].append(opinions)
        
        # 记录每个群体的平均观点
        for group_id in range(config["simulation"]["num_groups"]):
            group_agents = [agent for agent in state["agents"] if agent["group_id"] == group_id]
            if group_agents:
                avg_opinion = np.mean([agent["opinion"] for agent in group_agents])
                results["group_opinions"][group_id].append(avg_opinion)
                
    return results

def visualize_results(results: dict, config: dict):
    """可视化模拟结果"""
    # 创建保存目录
    os.makedirs(config["visualization"]["save_path"], exist_ok=True)
    
    # 1. 网络结构可视化
    plt.figure(figsize=(10, 8))
    G = nx.Graph()
    for i in range(len(results["opinions"][0])):
        G.add_node(i)
    for i in range(len(results["opinions"][0])):
        for j in range(i+1, len(results["opinions"][0])):
            if np.random.random() < config["network"]["density"]:
                G.add_edge(i, j)
    
    pos = nx.spring_layout(G)
    node_colors = [results["opinions"][-1][i] for i in range(len(results["opinions"][0]))]
    nodes = nx.draw_networkx_nodes(G, pos, node_color=node_colors, cmap=plt.cm.coolwarm, 
                                  node_size=100)
    nx.draw_networkx_edges(G, pos)
    plt.colorbar(nodes)
    plt.title("社交网络结构（节点颜色表示观点倾向）")
    plt.savefig(os.path.join(config["visualization"]["save_path"], "network_structure.png"), 
                bbox_inches='tight', dpi=config["visualization"]["dpi"])
    plt.close()
    
    # 2. 观点分布热图
    plt.figure(figsize=(12, 6))
    opinions_matrix = np.array(results["opinions"]).T
    im = plt.imshow(opinions_matrix, aspect='auto', cmap='coolwarm')
    plt.colorbar(im)
    plt.title("观点分布热图")
    plt.xlabel("时间步")
    plt.ylabel("智能体ID")
    plt.savefig(os.path.join(config["visualization"]["save_path"], "opinion_heatmap.png"), 
                bbox_inches='tight', dpi=config["visualization"]["dpi"])
    plt.close()
    
    # 3. 群体平均观点变化
    plt.figure(figsize=(10, 6))
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c']
    group_names = ['温和派', '积极派', '消极派']
    for group_id, group_opinions in enumerate(results["group_opinions"]):
        plt.plot(group_opinions, label=group_names[group_id % len(colors)], 
                color=colors[group_id % len(colors)], linewidth=2)
    plt.title("群体平均观点变化")
    plt.xlabel("轮次")
    plt.ylabel("平均观点值")
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.savefig(os.path.join(config["visualization"]["save_path"], "group_opinions.png"), 
                bbox_inches='tight', dpi=config["visualization"]["dpi"])
    plt.close()
    
    # 4. 沉默率变化
    plt.figure(figsize=(10, 6))
    silent_ratios = []
    for step_opinions in results["opinions"]:
        silent_ratio = sum(1 for x in step_opinions if abs(x) < 0.1) / len(step_opinions)
        silent_ratios.append(silent_ratio)
    plt.plot(silent_ratios)
    plt.title("沉默率变化")
    plt.xlabel("轮次")
    plt.ylabel("沉默率")
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.savefig(os.path.join(config["visualization"]["save_path"], "silent_ratio.png"), 
                bbox_inches='tight', dpi=config["visualization"]["dpi"])
    plt.close()
    
    # 5. 观点标准差变化
    plt.figure(figsize=(10, 6))
    std_opinions = [np.std(step_opinions) for step_opinions in results["opinions"]]
    plt.plot(std_opinions)
    plt.title("观点标准差变化")
    plt.xlabel("轮次")
    plt.ylabel("标准差")
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.savefig(os.path.join(config["visualization"]["save_path"], "opinion_std.png"), 
                bbox_inches='tight', dpi=config["visualization"]["dpi"])
    plt.close()
    
    # 6. 代表性Agent轨迹
    plt.figure(figsize=(10, 6))
    representative_agents = [0, len(results["opinions"][0])//2, -1]  # 选择三个代表性Agent
    for agent_id in representative_agents:
        trajectory = [step_opinions[agent_id] for step_opinions in results["opinions"]]
        plt.plot(trajectory, label=f"Agent {agent_id}")
    plt.title("代表性Agent观点轨迹")
    plt.xlabel("轮次")
    plt.ylabel("观点值")
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.savefig(os.path.join(config["visualization"]["save_path"], "agent_trajectories.png"), 
                bbox_inches='tight', dpi=config["visualization"]["dpi"])
    plt.close()

def get_config() -> Dict[str, Any]:
    """获取配置"""
    return {
        "num_steps": 100,
        "network_density": 0.1,
        "random_seed": 42,  # 添加随机种子
        "visualization": {
            "enabled": True,
            "save_path": "results/visualization"
        }
    }

def main():
    # 加载配置
    config = load_config("config/config.yaml")
    
    # 创建结果目录
    os.makedirs(os.path.dirname(config["output"]["results_path"]), exist_ok=True)
    os.makedirs(config["visualization"]["save_path"], exist_ok=True)
    
    # 运行模拟
    results = run_simulation(config)
    
    # 保存结果
    with open(config["output"]["results_path"], "w") as f:
        json.dump(results, f)
        
    # 可视化结果
    visualize_results(results, config)

if __name__ == "__main__":
    main() 